import os
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import rc as matplotlib_rc
from code.flaw import Flaw

# 设置中文字体
font = {'family': "YaHei Consolas Hybrid"}
matplotlib_rc("font", **font)


def get_all_flaw(path):
    """
    根据路径解析所有图片标签（xml文件）,获取所有瑕疵类
    :param path:待解析的图片路径
    :return: flaws 包含所有瑕疵的数组
    """
    file_names = os.listdir(path)
    # 获取所有类型
    flaws = []  # 所有flaw
    for file_name in file_names:
        if file_name[-3:] == "jpg":
            flaw = Flaw(pic_name=path + file_name)
            flaws.append(flaw)
    return flaws


def get_all_type_size(flaws, type_names, type_num):
    """
    根据传入的flaw数组，计算各种类型flaw的大小
    :param flaws:flaw数组
    :param type_names:对应的type_names数组
    :param type_nums:对应的type_nums数组
    :return: type_size各种类型的大小 [min,avg,max]
    """
    type_size = []  # 所有种类的大小   [min,avg,max]
    for i in range(len(type_names)):  # 初始化
        type_size.append([10000000, 0, 0])

    for flaw in flaws:
        if flaw.type == '正常': continue
        # 获取所有种类数量
        index = type_names.index(flaw.type)
        type_num[index] += 1
        # 获取所有种类的大小[min,sum,max]
        for size in flaw.get_size():
            type_size[index][1] += size
            if size < type_size[index][0]:
                type_size[index][0] = size
            if size > type_size[index][2]:
                type_size[index][2] = size
        type_size[index][1] /= len(list(flaw.get_size()))
    # 将大小的和转化为平均值
    for index in range(len(type_size)):
        type_size[index][1] /= type_num[index]
    return type_size


def draw_type_num(type_names, type_num, save_path=None):
    """
    根据type_names和type_num绘制柱状图，如果传入save_path，则保存图片
    :param type_names:作为x轴
    :param type_num:作为y轴
    :param save_path:保存路径
    :return:None
    """
    # 设置图片大小
    plt.figure(figsize=(10, 8))
    plt.bar(type_names, type_num)
    plt.grid(alpha=0.4, linestyle=':')
    plt.xticks(rotation=90)
    # 设置xy标签和图标题
    plt.xlabel("瑕疵类型")
    plt.ylabel("样本数量")
    plt.title("瑕疵类型-样本数量")
    # 保存图片
    if save_path:
        plt.savefig(save_path + "瑕疵类型-样本数量.svg")
    plt.show()


def draw_type_size(type_names, type_size_avg, type_size_min,
                   type_size_max, save_path=None):
    """
    根据type_names和type_num绘制柱状图，如果传入save_path，则保存图片
    :param type_size_avg: 平均值
    :param type_size_max: 最大值
    :param type_size_min: 最小值
    :param type_names:作为x轴
    :param save_path:保存路径
    :return:None
    """
    # 设置图片大小
    plt.figure(figsize=(10, 8))
    plt.bar(type_names, type_size_max, label="最大瑕疵")
    plt.bar(type_names, type_size_avg, label="平均瑕疵")
    plt.bar(type_names, type_size_min, label="最小瑕疵")
    plt.grid(alpha=0.4, linestyle=':')
    plt.xticks(rotation=90)
    plt.legend(loc="upper left")
    # 设置xy标签和图标题
    plt.xlabel("瑕疵类型")
    plt.ylabel("瑕疵大小")
    plt.title("瑕疵类型-瑕疵大小")
    # 保存图片
    if save_path:
        plt.savefig(save_path + "瑕疵类型-瑕疵大小.svg")
    plt.show()


def draw_flaw_pos(flaws, save_path=None, step=100):
    """
    画出所有瑕疵出现的位置
    :param flaws:所有瑕疵类数组
    :param save_path:保存路径
    :param step:区域宽度
    :return:None
    """
    # 定义等高线图的横纵坐标x，y
    x = range(0, 2560, step)
    y = range(0, 1920, step)
    # 将原始数据变成网格数据
    X, Y = np.meshgrid(x, y)

    # 设置图片大小
    plt.figure(figsize=(10, 8))
    # 各地点对应的高度数据
    height = np.zeros((len(y), len(x)))
    for flaw in flaws:
        if flaw.poses is None:
            continue
        for pos in flaw.poses:
            for i in range(pos['ymin'], pos['ymax'], step):
                for j in range(pos['xmin'], pos['xmax'], step):
                    height[(i // step) - 1][(j // step) - 1] += 1

    # 填充颜色
    plt.contourf(X, Y, height, 10, alpha=0.6, cmap=plt.cm.hot)
    # 绘制等高线
    C = plt.contour(X, Y, height, 10, colors='black')
    # 显示各等高线的数据标签
    plt.clabel(C, inline=True, fmt="%d", fontsize="x-large")
    # 保存图片
    if save_path:
        plt.savefig(save_path + "瑕疵坐标分布.svg")
    plt.show()


if __name__ == '__main__':
    # 获取所有瑕疵类
    flaws = get_all_flaw("../../data/train/")
    # 获取所有所有种类名称并去重
    type_names = list(set([flaw.type for flaw in flaws]))
    # 所有种类的数量
    type_num = [0] * len(type_names)

    # 删除无瑕疵数据
    index = type_names.index("正常")
    del type_names[index]
    del type_num[index]

    # 获取所有瑕疵种类的大小
    type_size = get_all_type_size(flaws, type_names, type_num)

    df = pd.DataFrame([type_names, type_num,
                       [size[0] for size in type_size],
                       [size[1] for size in type_size],
                       [size[2] for size in type_size]],
                      index=["类型", "数量", "最小瑕疵", "平均瑕疵", "最大瑕疵"])
    df_sorted = df.T.sort_values(by="数量", ascending=False)
    print(df_sorted)

    draw_type_num(df_sorted[0:]['类型'],
                  df_sorted[0:]['数量'], save_path="./")

    draw_type_size(df_sorted[0:]['类型'],
                   df_sorted[0:]['平均瑕疵'],
                   df_sorted[0:]['最小瑕疵'],
                   df_sorted[0:]['最大瑕疵'], save_path="./")

    draw_flaw_pos(flaws, save_path="./")
